#pragma once

#include <cmath>
#include <utility>

namespace jianzhi_055
{

struct TreeNode
{
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) { }
};

class Solution
{
public:
    bool isBalanced(TreeNode* root)
    {
        return is_balanced_with_depth(root).first;
    }

private:
    std::pair<bool, int> is_balanced_with_depth(TreeNode* root)
    {
        if (root == nullptr)
            return {true, 0};
        auto [l_balanced, l_depth] = is_balanced_with_depth(root->left);
        auto [r_balanced, r_depth] = is_balanced_with_depth(root->right);
        auto balanced = l_balanced && r_balanced && (std::abs(l_depth - r_depth) <= 1);
        auto depth = std::max(l_depth, r_depth) + 1;
        return {balanced, depth};
    }
};

}  /* namespace jianzhi_055 */
